import os
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
from scipy.signal import welch
from right_data_loader import all_data_loader

from dataProcessor import DataProcessor

from sklearn.cluster import KMeans
from sklearn.decomposition import PCA

def data_transformation(path):
    day0_data, day0_labels, day0_info = all_data_loader(path)
    day0_label = day0_info['success_cond']

    data_processor = DataProcessor()

    fixed_day0_data = []
    fixed_day0_label = []
    fixed_day0_labels = []

    max_len = 0
    min_len = int(2 * day0_info["rhd_fs"])

    # 1. 对数据进行加箱、平滑、标准化、转置、并调整数据大小
    for did, arr in enumerate(day0_data):
        # 持续时长超1s的trial舍弃
        if arr.shape[1] > int(1 * day0_info["rhd_fs"]):
            continue

        # data
        # 低通
        arr = data_processor.lowpass_filter(arr, highcut=200, fs=day0_info["rhd_fs"])
        # 带通滤波
        # arr = data_processor.bandpass_filter(arr, lowcut=70, highcut=200, fs=day0_info['rhd_fs'])
        # 带阻滤波
        arr = data_processor.notch_filter(arr, notched_fs=50, fs=day0_info["rhd_fs"], Q=0.2)
        # 降采样到一定长度
        # arr = data_processor.resampling(arr, 10000)
        # 采样率降采样
        arr = data_processor.down_sampling(arr, factor=20)

        # 标准化
        mean_value = np.mean(arr, axis=1, keepdims=True)
        std_value = np.std(arr, axis=1, keepdims=True)
        arr = (arr - mean_value) / std_value

        if arr.shape[1] > max_len:
            max_len = arr.shape[1]

        if arr.shape[1] < min_len:
            min_len = arr.shape[1]

        fixed_day0_data.append(arr)
        fixed_day0_labels.append(day0_labels[did].T)
        fixed_day0_label.append(day0_label[did] - 1)

    return fixed_day0_data, fixed_day0_labels, fixed_day0_label


def lfp_kmeans(path):
    data, kinemat, labels = data_transformation(path)

    flatten_data = [arr.flatten() for arr in data]
    X = np.vstack(flatten_data)

    model = KMeans(n_clusters=4, random_state=42)
    model.fit(X)
    pre_labels = model.labels_

    pca = PCA(n_components=2, random_state=42)
    X_pca = pca.fit_transform(X)

    n = 0
    for i in range(len(labels)):
        if labels[i] == pre_labels[i]:
            n += 1

    print('acc=', n / len(labels))

    true0 = 0
    true1 = 0
    true2 = 0
    true3 = 0

    pre0 = 0
    pre1 = 0
    pre2 = 0
    pre3 = 0

    for i in range(len(labels)):
        if labels[i] == 0:
            true0 += 1
        elif labels[i] == 1:
            true1 += 1
        elif labels[i] == 2:
            true2 += 1
        elif labels[i] == 3:
            true3 += 1

    print(f'True labels,0:{true0},1:{true1},2:{true2},3:{true3}')

    fig, (ax1, ax2) = plt.subplots(1, 2)
    for i in range(X_pca.shape[0]):
        if pre_labels[i] == 0:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='red')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='red')
            pre0 += 1
        elif pre_labels[i] == 1:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='blue')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='blue')
            pre1 += 1
        elif pre_labels[i] == 2:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='green')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='green')
            pre2 += 1
        elif pre_labels[i] == 3:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='yellow')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='yellow')
            pre3 += 1

    print(f'Predicted labels,0:{pre0},1:{pre1},2:{pre2},3:{pre3}')
    plt.show()


def psd_kmeans(path):
    data, kinemat, labels = data_transformation(path)

    all_psds = []
    frequencies = None

    for did, arr in enumerate(data):
        psds = []
        for k in range(arr.shape[0]):
            freqs, psd = welch(arr[k], fs=1000)
            psds.append(psd)
            if frequencies is None:
                frequencies = freqs
                print(frequencies)
        psds = np.array(psds)
        freq_mask = (frequencies >= 70) & (frequencies <= 120)
        psd_band = psds[:, freq_mask].mean(axis=1)
        # psd_band = psds[:, :].mean(axis=1)
        all_psds.append(psd_band)

    flatten_data = [arr.flatten() for arr in all_psds]
    X = np.vstack(flatten_data)

    model = KMeans(n_clusters=4, random_state=42)
    model.fit(X)
    pre_labels = model.labels_

    pca = PCA(n_components=2, random_state=42)
    X_pca = pca.fit_transform(X)

    n = 0
    for i in range(len(labels)):
        if labels[i] == pre_labels[i]:
            n += 1

    print('acc=', n / len(labels))

    true0 = 0
    true1 = 0
    true2 = 0
    true3 = 0

    pre0 = 0
    pre1 = 0
    pre2 = 0
    pre3 = 0

    red0 = 0
    red1 = 0
    red2 = 0
    red3 = 0

    blue0 = 0
    blue1 = 0
    blue2 = 0
    blue3 = 0

    green0 = 0
    green1 = 0
    green2 = 0
    green3 = 0

    yellow0 = 0
    yellow1 = 0
    yellow2 = 0
    yellow3 = 0

    for i in range(len(labels)):
        if labels[i] == 0:
            true0 += 1
        elif labels[i] == 1:
            true1 += 1
        elif labels[i] == 2:
            true2 += 1
        elif labels[i] == 3:
            true3 += 1
    print(f'True labels,0:{true0},1:{true1},2:{true2},3:{true3}')

    fig, (ax1, ax2) = plt.subplots(1, 2)
    for i in range(X_pca.shape[0]):
        if pre_labels[i] == 0:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='red')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='red')
            if labels[i] == 0:
                red0 += 1
            elif labels[i] == 1:
                red1 += 1
            elif labels[i] == 2:
                red2 += 1
            elif labels[i] == 3:
                red3 += 1
            pre0 += 1
        elif pre_labels[i] == 1:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='blue')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='blue')
            if labels[i] == 0:
                blue0 += 1
            elif labels[i] == 1:
                blue1 += 1
            elif labels[i] == 2:
                blue2 += 1
            elif labels[i] == 3:
                blue3 += 1
            pre1 += 1
        elif pre_labels[i] == 2:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='green')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='green')
            if labels[i] == 0:
                green0 += 1
            elif labels[i] == 1:
                green1 += 1
            elif labels[i] == 2:
                green2 += 1
            elif labels[i] == 3:
                green3 += 1
            pre2 += 1
        elif pre_labels[i] == 3:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='yellow')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='yellow')
            if labels[i] == 0:
                yellow0 += 1
            elif labels[i] == 1:
                yellow1 += 1
            elif labels[i] == 2:
                yellow2 += 1
            elif labels[i] == 3:
                yellow3 += 1
            pre3 += 1

    print(f'Predicted labels,0:{pre0},1:{pre1},2:{pre2},3:{pre3}')
    print(f'red,0:{red0},1:{red1},2:{red2},3:{red3}')
    print(f'blue,0:{blue0},1:{blue1},2:{blue2},3:{blue3}')
    print(f'green,0:{green0},1:{green1},2:{green2},3:{green3}')
    print(f'yellow,0:{yellow0},1:{yellow1},2:{yellow2},3:{yellow3}')

    plt.show()

def win_psd_kmeans(path, num_win, win_pos):
    data, kinemat, labels = data_transformation(path)

    data_len = 10000
    for i in range(len(data)):
        movement_len = data[i].shape[1]
        if movement_len < data_len:
            data_len = movement_len

    print('min_len', data_len)

    win_len = int(np.ceil(data_len // num_win))

    for i in range(len(data)):
        data[i] = data[i][:, :data_len]
        data[i] = data[i][:, win_len * (win_pos - 1):win_len * win_pos]

    all_psds = []
    frequencies = None

    for i, arr in enumerate(data):
        psds = []
        for k in range(arr.shape[0]):
            freqs, psd = welch(arr[k], fs=1000)
            psds.append(psd)
            if frequencies is None:
                frequencies = freqs
                print(frequencies)
        psds = np.array(psds)
        freq_mask = (frequencies >= 70) & (frequencies <= 120)
        psd_band = psds[:, freq_mask].mean(axis=1)
        # psd_band = psds[:, :].mean(axis=1)
        all_psds.append(psd_band)

    flatten_data = [arr.flatten() for arr in all_psds]
    X = np.vstack(flatten_data)

    model = KMeans(n_clusters=4, random_state=42)
    model.fit(X)
    pre_labels = model.labels_

    pca = PCA(n_components=2, random_state=42)
    X_pca = pca.fit_transform(X)

    n = 0
    for i in range(len(labels)):
        if labels[i] == pre_labels[i]:
            n += 1

    print('acc=', n / len(labels))

    true0 = 0
    true1 = 0
    true2 = 0
    true3 = 0

    pre0 = 0
    pre1 = 0
    pre2 = 0
    pre3 = 0

    red0 = 0
    red1 = 0
    red2 = 0
    red3 = 0

    blue0 = 0
    blue1 = 0
    blue2 = 0
    blue3 = 0

    green0 = 0
    green1 = 0
    green2 = 0
    green3 = 0

    yellow0 = 0
    yellow1 = 0
    yellow2 = 0
    yellow3 = 0

    for i in range(len(labels)):
        if labels[i] == 0:
            true0 += 1
        elif labels[i] == 1:
            true1 += 1
        elif labels[i] == 2:
            true2 += 1
        elif labels[i] == 3:
            true3 += 1

    print(f'True labels,0:{true0},1:{true1},2:{true2},3:{true3}')

    fig, (ax1, ax2) = plt.subplots(1, 2)
    for i in range(X_pca.shape[0]):
        if pre_labels[i] == 0:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='red')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='red')
            if labels[i] == 0:
                red0 += 1
            elif labels[i] == 1:
                red1 += 1
            elif labels[i] == 2:
                red2 += 1
            elif labels[i] == 3:
                red3 += 1
            pre0 += 1
        elif pre_labels[i] == 1:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='blue')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='blue')
            if labels[i] == 0:
                blue0 += 1
            elif labels[i] == 1:
                blue1 += 1
            elif labels[i] == 2:
                blue2 += 1
            elif labels[i] == 3:
                blue3 += 1
            pre1 += 1
        elif pre_labels[i] == 2:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='green')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='green')
            if labels[i] == 0:
                green0 += 1
            elif labels[i] == 1:
                green1 += 1
            elif labels[i] == 2:
                green2 += 1
            elif labels[i] == 3:
                green3 += 1
            pre2 += 1
        elif pre_labels[i] == 3:
            ax1.scatter(X_pca[i, 0], X_pca[i, 1], color='yellow')
            ax2.plot(kinemat[i][0, :], kinemat[i][1, :], color='yellow')
            if labels[i] == 0:
                yellow0 += 1
            elif labels[i] == 1:
                yellow1 += 1
            elif labels[i] == 2:
                yellow2 += 1
            elif labels[i] == 3:
                yellow3 += 1
            pre3 += 1

    print(f'Predicted labels,0:{pre0},1:{pre1},2:{pre2},3:{pre3}')
    print(f'red,0:{red0},1:{red1},2:{red2},3:{red3}')
    print(f'blue,0:{blue0},1:{blue1},2:{blue2},3:{blue3}')
    print(f'green,0:{green0},1:{green1},2:{green2},3:{green3}')
    print(f'yellow,0:{yellow0},1:{yellow1},2:{yellow2},3:{yellow3}')

    plt.show()


if __name__ == '__main__':

    path = r'D:\data\data_save_raw_right\rm08\20240229'
    # data_transformation(path)
    # lfp_kmeans(path)
    psd_kmeans(path)

    num_win = 4
    win_pos = 1
    for i in range(num_win):
        i = i + 1
        win_psd_kmeans(path, num_win=num_win, win_pos=i)
